from datetime import datetime
import akshare as ak
import matplotlib.pyplot as plt
import backtrader as bt
import pandas as pd

"""
https://zhuanlan.zhihu.com/p/122183963
"""


class MyStrategy(bt.Strategy):
    params = (('maperiod', 20),)  # 全局设定交易策略的参数

    def __init__(self):
        self.dataclose = self.datas[0].close  # 指定价格序列
        # self.data_close = self.data.lines.close
        # 初始化交易指令，买卖价格和手续费
        self.order = None
        self.buy_price = None
        self.buy_comm = None

        # 添加移动均线指标
        self.sma = bt.indicators.SimpleMovingAverage(self.datas[0], period=self.params.maperiod)

    def log(self, txt, dt=None):
        ''' Logging function fot this strategy'''
        dt = dt or self.datas[0].datetime.date(0)  # 当初的日期
        print('%s, %s' % (dt.isoformat(), txt))

    def next(self):
        if self.order:  # 查看是否有指令执行，如果有则不执行这bar
            return

        if not self.position:  # 没有持仓
            if self.dataclose[0] > self.sma[0]:  # 执行买入条件判断：收盘价上涨突破20日均线
                self.order = self.buy()  # 执行买入
        else:
            if self.dataclose[0] < self.sma[0]:  # 执行卖出条件判断：收盘价跌破20日均线
                self.orde = self.sell()  # 执行卖出

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            # Buy/Sell order submitted/accepted to/by broker - Nothing to do
            return

        # Check if an order has been completed
        # Attention: broker could reject order if not enough cash
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(
                    'BUY EXECUTED, Price: %.2f, Cost: %.2f, Comm %.2f' %
                    (order.executed.price,
                     order.executed.value,
                     order.executed.comm))

                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
                self.log('买入后以后资产：%.2f 元' % self.broker.getvalue())
            else:  # Sell
                self.log('SELL EXECUTED, Price: %.2f, Cost: %.2f, Comm %.2f' %
                         (order.executed.price,
                          order.executed.value,
                          order.executed.comm))
                self.log('卖出后以后资产：%.2f 元' % self.broker.getvalue())
            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')

        # Write down: no pending order
        self.order = None

    def notify_trade(self, trade):  # 交易执行后，在这里处理
        if not trade.isclosed:
            return

        self.log('OPERATION PROFIT, GROSS %.2f, NET %.2f' % (trade.pnl, trade.pnlcomm))  # 记录下盈利数据。


def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    # df.date = pd.to_datetime(df.date, format='%Y-%m-%d', utc=True)
    # df.set_index('date', inplace=True)
    df.index = pd.to_datetime(df.date)
    # df = df.sort_index()

    return df


if __name__ == '__main__':
    code = 'sz000002'
    start_date = datetime(2022, 1, 1)
    end_date = datetime(2023, 3, 31)
    df = get_data(code)
    df = df.sort_index()
    cerebro = bt.Cerebro()
    data = bt.feeds.PandasData(dataname=df, fromdate=start_date, todate=end_date)  # 加载数据
    cerebro.addstrategy(MyStrategy)  # 将交易策略加载到回测系统中
    cerebro.adddata(data)
    start_cash = 300000
    # Add a FixedSize sizer according to the stake
    cerebro.addsizer(bt.sizers.FixedSize, stake=100)
    cerebro.broker.setcash(start_cash)  # 设置初始资本
    cerebro.broker.setcommission(commission=0.002)
    cerebro.run()

    port_value = cerebro.broker.getvalue()  # 获取回测结束后的总资金
    pnl = port_value - start_cash  # 盈亏统计

    print(f"初始基金： {start_cash}\n回测期间：{start_date.strftime('%Y%m%d')}:{end_date.strftime('%Y%m%d')}")
    print(f"总基金：{round(port_value, 2)}")
    print(f"净收益：{round(pnl, 2)}")

    # cerebro.plot(style='candlestick')  # 画图
    cerebro.plot()
